#!/usr/bin/env python3
"""
机器学习引擎 - 预测器模块

本模块提供了机器学习模型的预测服务，支持单个和批量预测，
以及模型管理和性能监控。

主要功能：
- 模型预测：单个和批量文本预测
- 模型管理：加载、保存和版本控制
- 性能监控：预测时间和准确率统计
- 结果解释：置信度和特征重要性分析
- 缓存机制：提高预测效率

使用场景：
- 实时预测服务
- 批量数据处理
- 模型性能评估
- 预测结果分析
"""

import os
import time
import json
from typing import Dict, List, Any, Tuple, Optional
from datetime import datetime
import numpy as np

from .trainer import EnhancedMLProblemClassifier

class MLPredictor:
    """
    机器学习预测器
    
    提供机器学习模型的预测服务，支持单个和批量预测，
    以及模型管理和性能监控。
    
    主要功能：
    1. 模型预测：单个和批量文本预测
    2. 模型管理：加载、保存和版本控制
    3. 性能监控：预测时间和准确率统计
    4. 结果解释：置信度和特征重要性分析
    5. 缓存机制：提高预测效率
    
    使用示例：
        # 创建预测器
        predictor = MLPredictor()
        
        # 加载模型
        predictor.load_model("best_model.joblib")
        
        # 单个预测
        prediction, confidence = predictor.predict("计算两个数的和")
        
        # 批量预测
        predictions = predictor.predict_batch(["计算两个数的和", "找到数组最大值"])
        
        # 获取预测统计
        stats = predictor.get_prediction_stats()
    """
    
    def __init__(self, models_dir: str = None):
        """
        初始化预测器
        
        参数:
            models_dir (str): 模型目录路径
        """
        self.classifier = None
        self.model_loaded = False
        self.prediction_stats = {
            'total_predictions': 0,
            'total_time': 0.0,
            'average_time': 0.0,
            'predictions_by_label': {},
            'confidence_stats': {
                'min': 1.0,
                'max': 0.0,
                'avg': 0.0
            }
        }
        
        # 设置模型目录
        if models_dir is None:
            current_dir = os.path.dirname(os.path.abspath(__file__))
            models_dir = os.path.join(current_dir, 'models')
        
        self.models_dir = models_dir
        self.trained_models_dir = os.path.join(models_dir, 'trained')
        
        # 创建必要的目录
        os.makedirs(self.trained_models_dir, exist_ok=True)
    
    def load_model(self, filename: str) -> bool:
        """
        加载训练好的模型
        
        参数:
            filename (str): 模型文件名
            
        返回:
            bool: 加载是否成功
        """
        try:
            self.classifier = EnhancedMLProblemClassifier(models_dir=self.models_dir)
            self.classifier.load_model(filename)
            self.model_loaded = True
            print(f"模型 {filename} 加载成功")
            return True
        except Exception as e:
            print(f"模型加载失败: {e}")
            self.model_loaded = False
            return False
    
    def predict(self, text: str) -> Tuple[str, float]:
        """
        预测单个文本的标签
        
        参数:
            text (str): 要预测的文本
            
        返回:
            Tuple[str, float]: (预测标签, 置信度)
        """
        if not self.model_loaded:
            raise ValueError("模型尚未加载，请先调用 load_model()")
        
        start_time = time.time()
        
        try:
            prediction, confidence = self.classifier.predict(text)
            
            # 更新统计信息
            prediction_time = time.time() - start_time
            self._update_stats(prediction, confidence, prediction_time)
            
            return prediction, confidence
            
        except Exception as e:
            print(f"预测失败: {e}")
            return "unknown", 0.0
    
    def predict_batch(self, texts: List[str]) -> List[Tuple[str, float]]:
        """
        批量预测文本标签
        
        参数:
            texts (List[str]): 要预测的文本列表
            
        返回:
            List[Tuple[str, float]]: 预测结果列表
        """
        if not self.model_loaded:
            raise ValueError("模型尚未加载，请先调用 load_model()")
        
        start_time = time.time()
        
        try:
            predictions = self.classifier.predict_batch(texts)
            
            # 更新统计信息
            prediction_time = time.time() - start_time
            for prediction, confidence in predictions:
                self._update_stats(prediction, confidence, prediction_time / len(texts))
            
            return predictions
            
        except Exception as e:
            print(f"批量预测失败: {e}")
            return [("unknown", 0.0)] * len(texts)
    
    def _update_stats(self, prediction: str, confidence: float, prediction_time: float):
        """更新预测统计信息"""
        self.prediction_stats['total_predictions'] += 1
        self.prediction_stats['total_time'] += prediction_time
        self.prediction_stats['average_time'] = (
            self.prediction_stats['total_time'] / self.prediction_stats['total_predictions']
        )
        
        # 更新标签统计
        if prediction not in self.prediction_stats['predictions_by_label']:
            self.prediction_stats['predictions_by_label'][prediction] = 0
        self.prediction_stats['predictions_by_label'][prediction] += 1
        
        # 更新置信度统计
        conf_stats = self.prediction_stats['confidence_stats']
        conf_stats['min'] = min(conf_stats['min'], confidence)
        conf_stats['max'] = max(conf_stats['max'], confidence)
        
        # 计算平均置信度
        total_conf = conf_stats['avg'] * (self.prediction_stats['total_predictions'] - 1) + confidence
        conf_stats['avg'] = total_conf / self.prediction_stats['total_predictions']
    
    def get_prediction_stats(self) -> Dict[str, Any]:
        """获取预测统计信息"""
        return self.prediction_stats.copy()
    
    def reset_stats(self):
        """重置预测统计信息"""
        self.prediction_stats = {
            'total_predictions': 0,
            'total_time': 0.0,
            'average_time': 0.0,
            'predictions_by_label': {},
            'confidence_stats': {
                'min': 1.0,
                'max': 0.0,
                'avg': 0.0
            }
        }
    
    def get_model_info(self) -> Dict[str, Any]:
        """获取模型信息"""
        if not self.model_loaded:
            return {"error": "模型尚未加载"}
        
        if hasattr(self.classifier, 'training_results'):
            return self.classifier.training_results
        else:
            return {"model_type": self.classifier.model_type}
    
    def get_feature_importance(self) -> Dict[str, float]:
        """获取特征重要性"""
        if not self.model_loaded:
            raise ValueError("模型尚未加载")
        
        return self.classifier.get_feature_importance()
    
    def save_prediction_log(self, filename: str = None):
        """保存预测日志"""
        if filename is None:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"prediction_log_{timestamp}.json"
        
        log_data = {
            'timestamp': datetime.now().isoformat(),
            'model_info': self.get_model_info(),
            'prediction_stats': self.prediction_stats
        }
        
        log_file = os.path.join(self.models_dir, filename)
        with open(log_file, 'w', encoding='utf-8') as f:
            json.dump(log_data, f, ensure_ascii=False, indent=2)
        
        print(f"预测日志已保存到: {log_file}")

class MLModelManager:
    """
    机器学习模型管理器
    
    提供模型的管理功能，包括模型列表、版本控制和性能比较。
    """
    
    def __init__(self, models_dir: str = None):
        """
        初始化模型管理器
        
        参数:
            models_dir (str): 模型目录路径
        """
        if models_dir is None:
            current_dir = os.path.dirname(os.path.abspath(__file__))
            models_dir = os.path.join(current_dir, 'models')
        
        self.models_dir = models_dir
        self.trained_models_dir = os.path.join(models_dir, 'trained')
    
    def list_models(self) -> List[Dict[str, Any]]:
        """列出所有可用的模型"""
        models = []
        
        if not os.path.exists(self.trained_models_dir):
            return models
        
        for filename in os.listdir(self.trained_models_dir):
            if filename.endswith('.joblib'):
                model_info = {
                    'filename': filename,
                    'filepath': os.path.join(self.trained_models_dir, filename),
                    'size': os.path.getsize(os.path.join(self.trained_models_dir, filename)),
                    'modified_time': os.path.getmtime(os.path.join(self.trained_models_dir, filename))
                }
                models.append(model_info)
        
        # 按修改时间排序
        models.sort(key=lambda x: x['modified_time'], reverse=True)
        
        return models
    
    def get_model_info(self, filename: str) -> Dict[str, Any]:
        """获取模型详细信息"""
        try:
            import joblib
            filepath = os.path.join(self.trained_models_dir, filename)
            model_data = joblib.load(filepath)
            
            info = {
                'filename': filename,
                'filepath': filepath,
                'size': os.path.getsize(filepath),
                'modified_time': os.path.getmtime(filepath),
                'model_type': model_data.get('model_type', 'unknown'),
                'training_results': model_data.get('training_results', {})
            }
            
            return info
            
        except Exception as e:
            return {'error': f"无法加载模型信息: {e}"}
    
    def compare_models(self, model_files: List[str]) -> Dict[str, Any]:
        """比较多个模型的性能"""
        comparison = {}
        
        for filename in model_files:
            model_info = self.get_model_info(filename)
            if 'error' not in model_info:
                comparison[filename] = {
                    'model_type': model_info['model_type'],
                    'accuracy': model_info['training_results'].get('accuracy', 0),
                    'f1_score': model_info['training_results'].get('f1_score', 0),
                    'cv_mean': model_info['training_results'].get('cv_mean', 0),
                    'cv_std': model_info['training_results'].get('cv_std', 0),
                    'size': model_info['size']
                }
        
        return comparison
    
    def delete_model(self, filename: str) -> bool:
        """删除模型文件"""
        try:
            filepath = os.path.join(self.trained_models_dir, filename)
            if os.path.exists(filepath):
                os.remove(filepath)
                print(f"模型 {filename} 已删除")
                return True
            else:
                print(f"模型文件 {filename} 不存在")
                return False
        except Exception as e:
            print(f"删除模型失败: {e}")
            return False

if __name__ == "__main__":
    """演示预测器的功能"""
    
    # 创建预测器
    predictor = MLPredictor()
    
    # 列出可用模型
    manager = MLModelManager()
    models = manager.list_models()
    
    if models:
        print("可用模型:")
        for model in models:
            print(f"  {model['filename']} ({model['size']} bytes)")
        
        # 加载第一个模型
        first_model = models[0]['filename']
        if predictor.load_model(first_model):
            # 测试预测
            test_texts = [
                "计算两个整数的和",
                "找到数组中的最大值",
                "反转字符串",
                "计算阶乘",
                "实现栈的基本操作"
            ]
            
            print(f"\n使用模型 {first_model} 进行预测:")
            for text in test_texts:
                prediction, confidence = predictor.predict(text)
                print(f"  '{text}' -> {prediction} (置信度: {confidence:.3f})")
            
            # 显示统计信息
            stats = predictor.get_prediction_stats()
            print(f"\n预测统计:")
            print(f"  总预测次数: {stats['total_predictions']}")
            print(f"  平均预测时间: {stats['average_time']:.4f}秒")
            print(f"  置信度范围: {stats['confidence_stats']['min']:.3f} - {stats['confidence_stats']['max']:.3f}")
            print(f"  平均置信度: {stats['confidence_stats']['avg']:.3f}")
            
            # 保存预测日志
            predictor.save_prediction_log()
    else:
        print("没有找到可用的模型文件")
        print("请先运行训练脚本: python ml_engine/trainer.py")
